#!/usr/bin/env python3
"""
Go Binary Analyzer for Tongyi Lingma
Searches for the 245-byte mapping table and custom Base64 alphabet
"""

import struct
import sys
import os
import re

def analyze_go_binary(binary_path):
    """Analyze Go binary for encryption-related data"""
    
    print(f"Analyzing {binary_path}...")
    
    # Known mapping values from our analysis
    known_mappings = {
        0xa1: 0x6d,  # m
        0xf3: 0x6f,  # o  
        0x6b: 0x64,  # d
        0x31: 0x65,  # e
        0xb0: 0x6c,  # l
        0x69: 0x22,  # "
        0x01: 0x3a,  # :
        0x90: 0x20   # space
    }
    
    # Custom Base64 alphabet we know
    custom_base64 = b"!#%&()*,.@ABCDEFGHIJKLMNOPQRSTUVWXYZ^_abcdefghijklmnopqrstuvwxyz"
    
    with open(binary_path, 'rb') as f:
        data = f.read()
    
    print(f"Binary size: {len(data)} bytes")
    
    # 1. Search for custom Base64 alphabet
    print("\n=== Searching for Custom Base64 Alphabet ===")
    base64_pos = data.find(custom_base64)
    if base64_pos != -1:
        print(f"✅ Found custom Base64 alphabet at offset 0x{base64_pos:x}")
        # Show some context
        start = max(0, base64_pos - 50)
        end = min(len(data), base64_pos + 100)
        context = data[start:end]
        print(f"Context: {context}")
    else:
        print("❌ Custom Base64 alphabet not found as continuous string")
        
        # Try to find parts of it
        parts = [b"!#%&()", b"ABCDEFGHIJKLM", b"abcdefghijklm", b"^_"]
        for part in parts:
            pos = data.find(part)
            if pos != -1:
                print(f"Found part '{part.decode()}' at 0x{pos:x}")
    
    # 2. Search for 256-byte arrays that could be mapping tables
    print("\n=== Searching for 256-byte Mapping Tables ===")
    candidates = []
    
    # Search for sequences that might be mapping tables
    # Look for 256 consecutive bytes aligned on reasonable boundaries
    for offset in range(0, len(data) - 256, 8):  # 8-byte alignment
        segment = data[offset:offset + 256]
        
        # Check if this could be our mapping table
        matches = 0
        for src, dst in known_mappings.items():
            if src < 256 and segment[src] == dst:
                matches += 1
        
        if matches >= 6:  # At least 6 matches
            candidates.append((offset, segment, matches))
            print(f"✅ Found mapping table candidate at 0x{offset:x} ({matches}/8 matches)")
    
    # 3. Search for Go string literals that contain encryption functions
    print("\n=== Searching for Go Function Names ===")
    encrypt_functions = [
        b"cosy/util/encrypt.CustomDecryptParts",
        b"cosy/util/encrypt.CustomDecryptV1", 
        b"cosy/util/encrypt.EncodeContext",
        b"cosy/util/encrypt.Base64Encode",
        b"cosy/util/encrypt.Base64Decode",
        b"cosy/util/encrypt.newEncoding",
        b"cosy/util/encrypt.shuffle"
    ]
    
    for func_name in encrypt_functions:
        pos = data.find(func_name)
        if pos != -1:
            print(f"Found function: {func_name.decode()} at 0x{pos:x}")
    
    # 4. Look for patterns that might indicate the mapping table
    print("\n=== Advanced Pattern Search ===")
    
    # Look for arrays where every element is unique (like a permutation table)
    for offset in range(0, len(data) - 256, 16):
        segment = data[offset:offset + 256]
        
        # Count unique bytes
        unique_count = len(set(segment))
        if unique_count >= 200:  # Most bytes are unique
            print(f"Found high-uniqueness array at 0x{offset:x} (unique: {unique_count}/256)")
            
            # Check known mappings
            matches = sum(1 for src, dst in known_mappings.items() 
                         if src < 256 and segment[src] == dst)
            if matches > 0:
                candidates.append((offset, segment, matches))
                print(f"  ✅ This array matches {matches}/8 known mappings!")
    
    # 5. Analyze the best candidates
    if candidates:
        print(f"\n=== Analyzing {len(candidates)} Candidates ===")
        
        # Sort by number of matches
        candidates.sort(key=lambda x: x[2], reverse=True)
        
        for i, (offset, segment, matches) in enumerate(candidates[:3]):
            print(f"\nCandidate #{i+1} at 0x{offset:x} ({matches} matches):")
            
            # Save the mapping table
            filename = f"mapping_candidate_{i+1}_0x{offset:x}.bin"
            with open(filename, 'wb') as f:
                f.write(segment)
            print(f"Saved to {filename}")
            
            # Test decryption with known encrypted data
            test_decryption(segment, filename.replace('.bin', '_test.txt'))
    
    return candidates

def test_decryption(mapping_table, output_file):
    """Test decryption with the found mapping table"""
    
    # Try to load our known encrypted data
    encrypted_files = ['decoded_result.bin', 'partial_decrypt.bin']
    
    for encrypted_file in encrypted_files:
        if os.path.exists(encrypted_file):
            print(f"Testing decryption with {encrypted_file}...")
            
            with open(encrypted_file, 'rb') as f:
                encrypted_data = f.read()
            
            # Apply the mapping table
            try:
                decrypted = bytes(mapping_table[b] for b in encrypted_data if b < 256)
                
                # Check if result looks like JSON
                if b'{"' in decrypted[:100] or b'"questionText"' in decrypted:
                    print(f"✅ Successful decryption!")
                    
                    with open(output_file, 'wb') as f:
                        f.write(decrypted)
                    print(f"Saved decrypted result to {output_file}")
                    
                    # Show preview
                    preview = decrypted[:500].decode('utf-8', errors='ignore')
                    print(f"Preview: {preview[:200]}...")
                    return True
                else:
                    print("❌ Decryption result doesn't look like valid JSON")
                    
            except Exception as e:
                print(f"❌ Decryption failed: {e}")
    
    return False

def extract_go_symbols(binary_path):
    """Extract Go symbols and function information"""
    
    print("\n=== Extracting Go Symbol Information ===")
    
    with open(binary_path, 'rb') as f:
        data = f.read()
    
    # Look for Go build info
    build_info_marker = b"\xff Go buildinf:"
    pos = data.find(build_info_marker)
    if pos != -1:
        print(f"Found Go build info at 0x{pos:x}")
        # Try to extract build info
        info_start = pos + len(build_info_marker)
        info_data = data[info_start:info_start + 100]
        print(f"Build info: {info_data}")
    
    # Look for Go module info
    module_patterns = [
        b"cosy/util/encrypt",
        b"main.main",
        b"runtime.",
        b"go.mod"
    ]
    
    for pattern in module_patterns:
        positions = []
        start = 0
        while True:
            pos = data.find(pattern, start)
            if pos == -1:
                break
            positions.append(pos)
            start = pos + 1
            if len(positions) > 20:  # Limit output
                break
        
        if positions:
            print(f"Found '{pattern.decode()}' at {len(positions)} locations: {[hex(p) for p in positions[:5]]}")

if __name__ == "__main__":
    binary_path = "Lingma"
    
    if not os.path.exists(binary_path):
        print(f"Binary {binary_path} not found!")
        sys.exit(1)
    
    # Extract Go symbols first
    extract_go_symbols(binary_path)
    
    # Analyze binary for mapping tables
    candidates = analyze_go_binary(binary_path)
    
    if candidates:
        print(f"\n✅ Analysis complete! Found {len(candidates)} mapping table candidates.")
        print("Check the generated files for decryption results.")
    else:
        print("\n❌ No mapping table candidates found.")
        print("The mapping table might be generated at runtime or stored differently.")